"""
GIF renderer module for converting Three.js HTML animations to GIF format.
Handles browser automation and image processing.
"""

import asyncio
import time
import logging
import psutil
from typing import List, Tuple, Dict, Any
from pathlib import Path
from selenium import webdriver
from selenium.webdriver.chrome.service import Service
from selenium.webdriver.chrome.options import Options
from selenium.common.exceptions import WebDriverException
from webdriver_manager.chrome import ChromeDriverManager
from PIL import Image
from io import BytesIO
from tqdm import tqdm

class ResourceMonitor:
    """监控系统资源使用情况"""
    
    @staticmethod
    def check_resources() -> Dict[str, float]:
        memory = psutil.virtual_memory()
        cpu_percent = psutil.cpu_percent(interval=1)
        return {
            'memory_percent': memory.percent,
            'memory_available': memory.available / (1024 * 1024),  # MB
            'cpu_percent': cpu_percent
        }
    
    @staticmethod
    def is_resources_critical() -> bool:
        resources = ResourceMonitor.check_resources()
        return (resources['memory_percent'] > 90 or 
                resources['memory_available'] < 500 or  # 小于 500MB 可用内存
                resources['cpu_percent'] > 80)

class WebGLChecker:
    """检查 WebGL 支持情况"""
    
    @staticmethod
    def inject_webgl_check(driver: webdriver.Chrome) -> Dict[str, Any]:
        check_script = """
        return (function() {
            const canvas = document.createElement('canvas');
            let gl;
            try {
                gl = canvas.getContext('webgl') || canvas.getContext('experimental-webgl');
                return {
                    supported: !!gl,
                    vendor: gl ? gl.getParameter(gl.VENDOR) : null,
                    renderer: gl ? gl.getParameter(gl.RENDERER) : null,
                    version: gl ? gl.getParameter(gl.VERSION) : null,
                    shadingLanguageVersion: gl ? gl.getParameter(gl.SHADING_LANGUAGE_VERSION) : null
                };
            } catch (e) {
                return {
                    supported: false,
                    error: e.toString()
                };
            }
        })();
        """
        return driver.execute_script(check_script)

def initialize_webdriver(config: Dict[str, Any]) -> webdriver.Chrome:
    """Initialize the Chrome WebDriver with optimized settings."""
    try:
        logging.info("Setting up Chrome options...")
        options = webdriver.ChromeOptions()
        
        # 基础设置
        options.add_argument("--headless=new")
        options.add_argument("--no-sandbox")
        options.add_argument("--disable-dev-shm-usage")
        
        # WebGL 优化
        options.add_argument("--use-gl=desktop")
        options.add_argument("--enable-webgl")
        options.add_argument("--disable-software-rasterizer")
        
        # 性能优化
        options.add_argument("--disable-extensions")
        options.add_argument("--disable-logging")
        options.add_argument(f"--window-size={config['resolution'][0]},{config['resolution'][1]}")
        options.add_argument("--disable-notifications")
        
        # 实验性选项
        options.add_experimental_option("excludeSwitches", ["enable-logging"])
        options.add_experimental_option('w3c', True)
        
        # 设置页面加载策略
        options.page_load_strategy = 'eager'
        
        service = Service(
            executable_path=ChromeDriverManager().install(),
            log_path='chromedriver.log',
            service_args=['--verbose']
        )
        
        driver = webdriver.Chrome(service=service, options=options)
        
        # 增加超时设置
        driver.set_page_load_timeout(120)  # 2分钟
        driver.implicitly_wait(30)
        driver.set_script_timeout(60)
        
        # 验证 WebGL 支持
        webgl_info = WebGLChecker.inject_webgl_check(driver)
        if not webgl_info.get('supported'):
            raise RuntimeError(f"WebGL not supported: {webgl_info.get('error')}")
        
        logging.info(f"WebGL Info: {webgl_info}")
        return driver
        
    except Exception as e:
        logging.error(f"Failed to initialize WebDriver: {str(e)}", exc_info=True)
        raise

async def capture_frames_batch(
    driver: webdriver.Chrome,
    total_frames: int,
    resolution: Tuple[int, int],
    batch_size: int = 10
) -> List[Image.Image]:
    """分批捕获帧以避免内存问题"""
    all_frames = []
    batches = range(0, total_frames, batch_size)
    
    with tqdm(total=total_frames, desc="Capturing frames") as pbar:
        for start_frame in batches:
            end_frame = min(start_frame + batch_size, total_frames)
            batch_frames = []
            
            for frame in range(start_frame, end_frame):
                if ResourceMonitor.is_resources_critical():
                    logging.warning("System resources critical, waiting...")
                    await asyncio.sleep(5)
                    continue
                
                # 执行动画更新
                driver.execute_script(f"animate({frame});")
                
                # 捕获截图
                screenshot = driver.get_screenshot_as_png()
                image = Image.open(BytesIO(screenshot))
                image = image.resize(resolution)
                batch_frames.append(image)
                
                pbar.update(1)
            
            all_frames.extend(batch_frames)
            
            # 每批次后进行垃圾回收
            import gc
            gc.collect()
            await asyncio.sleep(0.1)
    
    return all_frames

async def generate_gif_async(
    html_file: str,
    output_gif: str,
    frames: int,
    resolution: Tuple[int, int],
    frame_delay: int,
    config: Dict[str, Any]
) -> None:
    """异步生成 GIF 文件"""
    driver = None
    try:
        driver = initialize_webdriver(config)
        
        # 加载 HTML 文件
        file_url = f"file://{html_file}"
        logging.info(f"Loading file: {file_url}")
        driver.get(file_url)
        
        # 等待场景初始化
        await asyncio.sleep(2)
        
        # 分批捕获帧
        frames = await capture_frames_batch(
            driver,
            frames,
            resolution,
            batch_size=10
        )
        
        # 保存 GIF
        logging.info("Saving GIF...")
        frames[0].save(
            output_gif,
            save_all=True,
            append_images=frames[1:],
            duration=frame_delay,
            loop=0,
            optimize=False
        )
        
        logging.info(f"GIF saved to {output_gif}")
        
    except Exception as e:
        logging.error(f"Error during GIF generation: {str(e)}", exc_info=True)
        raise
    finally:
        if driver:
            driver.quit()

def generate_gif(
    html_file: str,
    output_gif: str,
    frames: int,
    resolution: Tuple[int, int],
    frame_delay: int,
    config: Dict[str, Any]
) -> None:
    """同步包装器"""
    try:
        asyncio.run(generate_gif_async(
            html_file,
            output_gif,
            frames,
            resolution,
            frame_delay,
            config
        ))
    except Exception as e:
        logging.error(f"Failed to generate GIF: {str(e)}", exc_info=True)
        raise